from hogs.solvers.flux.flux_functions import FluxFunction, LxfFlux, \
     LinearSolverFlux

from hogs.grids.reconstruct1d import ReconstructionType, MUSCL

import numpy
from optparse import OptionGroup

##########################################################################
# Options for the coupled acoustics and advection problem
##########################################################################
class AcousticsWithAdvectionOptions(object):
    """Options for the coupled acoustics and advection problem"""

    def __init__(self, u0=0.25, K0=0.25, rho0=1.0):        
        self.defaults = dict( u0=u0, K0=K0, rho0=rho0 )

    def get_options(self, parser):
        """Set options for the flux function"""

        txt = """Coupled acoustics and advection Jacobian matrix options"""
        opt = OptionGroup(parser, txt)

        opt.add_option("--u0", action="store", type="float",
                       dest="u0", default=self.defaults["u0"],
                       help="The background velocity for advection")

        opt.add_option("--K0", action="store", type="float",
                       dest="K0", default=self.defaults["K0"],
                       help="The bulk modulus")

        opt.add_option("--rho0", action="store", type="float",
                       dest="rho0", default=self.defaults["rho0"],
                       help="Background density")

        return opt

    def set_options(self, fluxfunc, options):
        options = options or self.defaults

        u0 = options.get("u0")
        K0 = options.get("K0")
        rho0 = options.get("rho0")

        # check the options
        self._check_options(u0, K0, rho0)

        # setup the coefficient matrix for the flux function
        self._setup_matrix(fluxfunc, u0, K0, rho0)

    def _setup_matrix(self, fluxfunc, u0, K0, rho0):

        fluxfunc.u0 = u0
        fluxfunc.K0 = K0
        fluxfunc.rho0 = rho0            

        fluxfunc.A = A = numpy.array([ [u0, K0, 0.0],
                                       [1.0/rho0, u0, 0.0],
                                       [0.0, 0.0, u0] ])            

        fluxfunc.c0 = numpy.sqrt( fluxfunc.K0/fluxfunc.rho0 )

        fluxfunc.Z0 = Z0 = fluxfunc.c0 * fluxfunc.rho0

        # store the Eigenvectors and Eigenvalues for the system
        fluxfunc.r1 = r1 = numpy.array( [-Z0, 1, 0.] )
        fluxfunc.r2 = r2 = numpy.array( [0, 0, 1.0 ] )
        fluxfunc.r3 = r3 = numpy.array( [Z0, 1, 0.0] )

        fluxfunc.r = [r1, r2, r3]

        fluxfunc.lambda1 = fluxfunc.u0 - fluxfunc.c0
        fluxfunc.lambda2 = fluxfunc.u0
        fluxfunc.lambda3 = fluxfunc.u0 + fluxfunc.c0

        fluxfunc.lam = [fluxfunc.lambda1, fluxfunc.lambda2, fluxfunc.lambda3]

    def _check_options(self, u0, K0, rho0):
        """Sanity check"""

        assert K0 > 0, "Invalid K0 value %g"%K0
        assert rho0 > 0, "Invalid rho0 value %g"%rho0
        
        if rho0 < 1e-10:
            raise ValueError("rho0 value %g too small"%rho0)

        c0 = numpy.sqrt( K0/rho0 )
        if c0 < 0:
            raise ValueError("Sound speed %g goes negative!"%c0)            

###########################################################################
#`AcousticswithAdvectionFlux1D`
###########################################################################
class AcousticswithAdvectionFlux1D(LinearSolverFlux):
    r"""Flux function for the acoustics problem coupled with advection.

    The system to be solved is given by:

    .. math::

        \begin{bmatrix} p \\ u \\ \phi \end{bmatrix}_t +
        \begin{bmatrix} u_0 \,\, K_0 \,\, 0 \\ \frac{1}{\rho_0}\,\, u_0 \,\,0
        \\ 0\,\, 0\,\, u_0 \end{bmatrix}\,\, \begin{bmatrix} p \\ u \\ \phi
        \end{bmatrix}_x = \begin{bmatrix} 0 \\ 0 \\ 0 \end{bmatrix}


    The eigenvalues for this system are :math:`\lambda^1 = u_0 - c_0`,
    :math:`\lambda^2 = u_0` and :math:`\lambda^3 = u_0 + c_0`, with
    corresponding eigenvalues:

    .. math::

        r^1 = \begin{bmatrix} -Z_0 \\ 1 \\ 0 \end{bmatrix} \,\, ,r^2 =
        \begin{bmatrix} 0 \\ 0 \\ 1 \end{bmatrix}\,\, ,r^3 =
        \begin{bmatrix} Z_0 \\ 1 \\ 0 \end{bmatrix}

    The wave strengths for the Riemann problem are given by:

    .. math::

        \alpha^1 = \frac{1}{2Z_0}\left[ -(p_r - p_l) + Z_0(u_r - u_l) \right ]

        \alpha^2 = \phi_r - \phi_l

        \alpha^3 = \frac{1}{2Z_0}\left[ (p_r - p_l) + Z_0(u_r - u_l) \right ]

    *FLux*:

    The fluxes for this system is computed in the standard form:

    .. math::

        F_{i + \frac{1}{2}} = AQ_{i} + \sum_{p=1}^{m} (\lambda^p)^-
        W_{i+\frac{1}{2}}^p + \frac{1}{2}\sum_{p=1}^{m}\left(
        |\lambda^p| (1 - \frac{\Delta t}{\Delta
        x}|\lambda^p|)\right)\tilde{W}_{i + \frac{1}{2}^p}

    resulting in a second order scheme.

    """
    def __init__(self, u0=0.25, K0=0.25, rho0=1.0, limiter_type=0):
        """Constructor.

        Parameters:
        ------------

        A : numpy.ndarray
            The coefficient matrix for the linear system.

        """
        LinearSolverFlux.__init__(self, limiter_type)
        self.optclass = AcousticsWithAdvectionOptions(u0, K0, rho0)

        c0 = numpy.sqrt( K0/rho0 )
        self.Z0 = c0 * rho0

    def get_alpha(self, deltaq):
        """Return the wave strengths for the Riemann problem"""
        Z0 = self.Z0

        alpha1 = 0.5/Z0 * (-deltaq[0] + Z0*deltaq[1])
        alpha2 = deltaq[2]
        alpha3 = 0.5/Z0 * (deltaq[0] + Z0*deltaq[1])

        return [alpha1, alpha2, alpha3]
        
    def compute_first_order_flux(self, face, grid, dt):
        """Compute the flux for each face"""
        nvar = grid.nvar
        Z0 = self.Z0

        lam = self.lam
        r = self.r

        f = face.flux
        ql, qr = face.get_conserved_variables(i=0)

        deltaq = qr - ql

        alpha = self.get_alpha(deltaq)
        
        # compute the first order term
        _f = numpy.dot( self.A, ql )
        for i in range(nvar):
            _f += min( lam[i], 0 ) * alpha[i] * r[i]

        # store the flux for this face
        f[:] = _f

    def compute_high_order_flux(self, face, grid, dt):
        """Compute high order correction terms for the flux"""

        k = dt/grid.dx
        nvar = grid.nvar
        nb = grid.nb

        lam = self.lam; r = self.r

        # get the difference at face i + 1/2
        ql, qr = face.get_conserved_variables(i=0)
        deltaq = qr - ql
        alpha = self.get_alpha(deltaq)

        # get the difference at face i + 3/2
        ql, qr = face.get_conserved_variables(i=1)
        deltaqp = qr - ql
        alphap = self.get_alpha(deltaqp)

        # get the difference at face i - 1/2
        ql, qr = face.get_conserved_variables(i=-1)
        deltaqm = qr - ql
        alpham = self.get_alpha(deltaqm)

        flux = face.flux

        for i in range(nvar):
            _lam = lam[i]; modlam = max(_lam, 0) - min(_lam, 0)
            
            # compute the flow parameter
            if lam > 0:
                fparam = alpham[i]/alpha[i]
            else:
                fparam = alphap[i]/alpha[i]

            # compute the limited wave
            wbar = self.limiter_func.limit(fparam) * alpha[i] * r[i]

            # add the flux contribution
            flux[:] += 0.5 * modlam * (1 - k * modlam) * wbar
        
    def get_options(self, parser):
        """Set options for the flux function"""

        # get the options from the parent 
        _opt = LinearSolverFlux.get_options(self, parser)

        # get options for the Jacobian
        opt = self.optclass.get_options(parser)

        return opt, _opt

    def set_options(self, options):
        """Perform flux specific setup.

        This flux function provides options to set the coefficient
        matrix A.
        
        """
        # set the parent class's options
        LinearSolverFlux.set_options(self, options)

        # set the Jacobian matrix options
        self.optclass.set_options(self, options)

###########################################################################
#`AcousticswithAdvectionLxfFlux1D`
###########################################################################
class AcousticswithAdvectionLxfFlux1D(LxfFlux):
    r"""Lax Friedrichs flux for the coupled acoustics and advection problem:

    .. math::

        F_{i + \frac{1}{2}} = \frac{1}{2}\left( F(Q_i) + F(Q_{i + 1})
        \right ) - \frac{\Delta x}{2\Delta t}\left( F(Q_{i + 1}) -
        F(Q_{i}) \right)

    The Lax Friedrichs flux is highly dissipative for the first order
    Godunov's method. For higher order reconstructions of the
    function, the dissipative nature is less profound. We implement
    three reconstruction algorithms to generate the boundary
    extrapolated values.

     * MUSCL : A slope is chosen in each cell approximating the first
       derivative. This method is formally second order.

     * PPM : Piecewise Parabolic interpolation method of Colella and
       Woodward. The method is formally third order in space.

     * ENO : Essentially non oscillatory reconstruction of Harten et
       al. The ENO method is capable of arbitrary high order
       reconstructions.

    """
    def __init__(self, u0=0.25, K0=0.25, rho0=1.0, muscl=0, ppm=-1, eno=-1):
        """Constructor."""
        LxfFlux.__init__(self, muscl, ppm, eno)

        # create the jacobian options class
        self.optclass = AcousticsWithAdvectionOptions(u0, K0, rho0)

    def get_options(self, parser):
        """Get options for the AcousticswithAdvectionLxfFlux1D"""

        # get the options for the Jacobian matrix
        jacobian_opts = self.optclass.get_options(parser)

        # get the reconstruction options
        reconstruct_opts = LxfFlux.get_options(self, parser)
        
        return jacobian_opts, reconstruct_opts

    def set_options(self, options):
        """Set options for the AcousticswithAdvectionLxfFlux1D"""
        options = options or self.defaults

        # set the reconstruction opts
        LxfFlux.set_options(self, options)

        # set the Jacobian options
        self.optclass.set_options(self, options)

    def _compute_unstable_flux(self, ql, qr):
        return numpy.dot( self.A, 0.5 * (ql + qr) )

    def _compute_max_sound_speed(self):
        return 1.0
        

flux_types = [
    ("Second order Godunov (default) ", AcousticswithAdvectionFlux1D),

    ("Lax Friedrichs with reconstruction", AcousticswithAdvectionLxfFlux1D)

    ]
